#include "codegen.h"
#include <sstream>

CodeGenerator::CodeGenerator() : labelCounter(0), stackOffset(0) {}

std::string CodeGenerator::generateLabel(const std::string& prefix) {
    return prefix + "_" + std::to_string(labelCounter++);
}

void CodeGenerator::emit(const std::string& instruction) {
    output += "    " + instruction + "\n";
}

void CodeGenerator::emitLabel(const std::string& label) {
    output += label + ":\n";
}

std::string CodeGenerator::generate(const Program* program) {
    output = ".section .text\n";
    output += ".global _start\n\n";
    
    // Generate code for each declaration
    for (const auto& decl : program->declarations) {
        generateDeclaration(decl.get());
    }
    
    // Add a simple _start function if no main function exists
    if (output.find("main:") == std::string::npos) {
        emitLabel("_start");
        emit("mov $60, %rax");  // sys_exit
        emit("mov $0, %rdi");   // exit code
        emit("syscall");
    }
    
    return output;
}

void CodeGenerator::generateDeclaration(const Declaration* decl) {
    if (auto varDecl = dynamic_cast<const VariableDeclaration*>(decl)) {
        generateVariableDeclaration(varDecl);
    } else if (auto funcDecl = dynamic_cast<const FunctionDeclaration*>(decl)) {
        generateFunctionDeclaration(funcDecl);
    }
}

void CodeGenerator::generateVariableDeclaration(const VariableDeclaration* decl) {
    // For global variables, we'll store them in .data section
    // For local variables, we'll use stack offsets
    if (decl->initializer) {
        generateExpression(decl->initializer.get());
    }
}

void CodeGenerator::generateFunctionDeclaration(const FunctionDeclaration* decl) {
    emitLabel(decl->name);
    
    // Function prologue
    emit("push %rbp");
    emit("mov %rsp, %rbp");
    
    // Allocate space for local variables
    stackOffset = 0;
    variableOffsets.clear();
    
    // Process parameters (they're on the stack)
    for (size_t i = 0; i < decl->parameters.size(); ++i) {
        // Parameters are at positive offsets from rbp
        variableOffsets[decl->parameters[i].second] = 8 + (i * 8);
    }
    
    // Generate function body
    if (decl->body) {
        generateBlockStatement(decl->body.get());
    }
    
    // Function epilogue
    emit("mov %rbp, %rsp");
    emit("pop %rbp");
    emit("ret");
    emit(""); // Empty line for readability
}

void CodeGenerator::generateStatement(const Statement* stmt) {
    if (auto exprStmt = dynamic_cast<const ExpressionStatement*>(stmt)) {
        generateExpressionStatement(exprStmt);
    } else if (auto retStmt = dynamic_cast<const ReturnStatement*>(stmt)) {
        generateReturnStatement(retStmt);
    } else if (auto ifStmt = dynamic_cast<const IfStatement*>(stmt)) {
        generateIfStatement(ifStmt);
    } else if (auto whileStmt = dynamic_cast<const WhileStatement*>(stmt)) {
        generateWhileStatement(whileStmt);
    } else if (auto blockStmt = dynamic_cast<const BlockStatement*>(stmt)) {
        generateBlockStatement(blockStmt);
    }
}

void CodeGenerator::generateExpressionStatement(const ExpressionStatement* stmt) {
    if (stmt->expression) {
        generateExpression(stmt->expression.get());
    }
}

void CodeGenerator::generateReturnStatement(const ReturnStatement* stmt) {
    if (stmt->expression) {
        generateExpression(stmt->expression.get());
        // Return value should be in %rax
    }
    emit("mov %rbp, %rsp");
    emit("pop %rbp");
    emit("ret");
}

void CodeGenerator::generateIfStatement(const IfStatement* stmt) {
    std::string elseLabel = generateLabel("else");
    std::string endLabel = generateLabel("endif");
    
    // Generate condition
    generateExpression(stmt->condition.get());
    
    // Jump if false
    emit("cmp $0, %rax");
    emit("je " + elseLabel);
    
    // Generate then branch
    generateStatement(stmt->thenBranch.get());
    emit("jmp " + endLabel);
    
    // Else label
    emitLabel(elseLabel);
    if (stmt->elseBranch) {
        generateStatement(stmt->elseBranch.get());
    }
    
    // End label
    emitLabel(endLabel);
}

void CodeGenerator::generateWhileStatement(const WhileStatement* stmt) {
    std::string loopLabel = generateLabel("loop");
    std::string endLabel = generateLabel("endloop");
    
    // Loop start
    emitLabel(loopLabel);
    
    // Generate condition
    generateExpression(stmt->condition.get());
    
    // Jump if false
    emit("cmp $0, %rax");
    emit("je " + endLabel);
    
    // Generate body
    generateStatement(stmt->body.get());
    
    // Jump back to loop start
    emit("jmp " + loopLabel);
    
    // End label
    emitLabel(endLabel);
}

void CodeGenerator::generateBlockStatement(const BlockStatement* stmt) {
    for (const auto& stmt : stmt->statements) {
        generateStatement(stmt.get());
    }
}

void CodeGenerator::generateExpression(const Expression* expr) {
    if (auto binaryExpr = dynamic_cast<const BinaryExpression*>(expr)) {
        generateBinaryExpression(binaryExpr);
    } else if (auto unaryExpr = dynamic_cast<const UnaryExpression*>(expr)) {
        generateUnaryExpression(unaryExpr);
    } else if (auto funcCall = dynamic_cast<const FunctionCall*>(expr)) {
        generateFunctionCall(funcCall);
    } else {
        // Simple expression (literal or identifier)
        if (expr->type == TokenType::INTEGER_LITERAL) {
            emit("mov $" + expr->value + ", %rax");
        } else if (expr->type == TokenType::IDENTIFIER) {
            // Load variable from memory
            auto it = variableOffsets.find(expr->value);
            if (it != variableOffsets.end()) {
                emit("mov " + std::to_string(it->second) + "(%rbp), %rax");
            } else {
                // Global variable
                emit("mov " + expr->value + "(%rip), %rax");
            }
        }
    }
}

void CodeGenerator::generateBinaryExpression(const BinaryExpression* expr) {
    // Generate left operand
    generateExpression(expr->left.get());
    pushToStack();
    
    // Generate right operand
    generateExpression(expr->right.get());
    
    // Right operand is now in %rax, left is on stack
    emit("pop %rcx"); // Left operand to %rcx
    
    switch (expr->operatorType) {
        case TokenType::PLUS:
            emit("add %rcx, %rax");
            break;
        case TokenType::MINUS:
            emit("sub %rcx, %rax");
            break;
        case TokenType::MULTIPLY:
            emit("imul %rcx, %rax");
            break;
        case TokenType::DIVIDE:
            emit("mov %rax, %rdx");
            emit("mov %rcx, %rax");
            emit("idiv %rcx");
            break;
        case TokenType::EQUAL:
            emit("cmp %rcx, %rax");
            emit("sete %al");
            emit("movzx %al, %rax");
            break;
        case TokenType::NOT_EQUAL:
            emit("cmp %rcx, %rax");
            emit("setne %al");
            emit("movzx %al, %rax");
            break;
        case TokenType::LESS:
            emit("cmp %rcx, %rax");
            emit("setl %al");
            emit("movzx %al, %rax");
            break;
        case TokenType::LESS_EQUAL:
            emit("cmp %rcx, %rax");
            emit("setle %al");
            emit("movzx %al, %rax");
            break;
        case TokenType::GREATER:
            emit("cmp %rcx, %rax");
            emit("setg %al");
            emit("movzx %al, %rax");
            break;
        case TokenType::GREATER_EQUAL:
            emit("cmp %rcx, %rax");
            emit("setge %al");
            emit("movzx %al, %rax");
            break;
        case TokenType::ASSIGN:
            // Store right operand (in %rax) to left operand location
            emit("mov %rax, " + std::to_string(variableOffsets[expr->left->value]) + "(%rbp)");
            break;
        default:
            break;
    }
}

void CodeGenerator::generateUnaryExpression(const UnaryExpression* expr) {
    generateExpression(expr->operand.get());
    
    switch (expr->operatorType) {
        case TokenType::NOT:
            emit("cmp $0, %rax");
            emit("sete %al");
            emit("movzx %al, %rax");
            break;
        case TokenType::MINUS:
            emit("neg %rax");
            break;
        default:
            break;
    }
}

void CodeGenerator::generateFunctionCall(const FunctionCall* call) {
    // For simplicity, we'll assume the function exists
    // In a real compiler, we'd need to handle calling conventions properly
    
    // Push arguments (in reverse order for x86-64 calling convention)
    for (int i = call->arguments.size() - 1; i >= 0; --i) {
        generateExpression(call->arguments[i].get());
        pushToStack();
    }
    
    // Call function
    emit("call " + call->functionName);
    
    // Clean up arguments from stack
    if (!call->arguments.empty()) {
        emit("add $" + std::to_string(call->arguments.size() * 8) + ", %rsp");
    }
}

void CodeGenerator::pushToStack() {
    emit("push %rax");
    stackOffset += 8;
}

void CodeGenerator::popFromStack() {
    emit("pop %rax");
    stackOffset -= 8;
}

std::string CodeGenerator::getRegisterForType(TokenType type) {
    switch (type) {
        case TokenType::INT:
        case TokenType::CHAR:
            return "%eax";
        case TokenType::FLOAT:
        case TokenType::DOUBLE:
            return "%xmm0";
        default:
            return "%rax";
    }
}
